import asyncio

from langchain_core.tools import StructuredTool
from pydantic import BaseModel, Field

"""
使用：StructuredTool.from_function创建tools
"""

class CalculatorInput(BaseModel):
    a: int = Field(description="first number")
    b: int = Field(description="second number")

def multiply(a: int, b: int) -> int:
    """Multiply two numbers"""
    return a * b


async def aMultiply(a: int, b: int) -> int:
    """async Multiply two numbers"""
    return a * b


async def main() -> None:
    """
     func参数：指定一个同步函数，当你在同步上下文中调用工具时，它会使用这个同步函数来执行操作。
    coroutine参数：指定一个异步函数，当你在异步上下文中调用工具时，它会使用这个异步函数来执行操作。
    """
    calculate = StructuredTool.from_function(
        func=multiply,
        name="multiply",
        description="multiply two numbers",
        args_schema=CalculatorInput,
        return_direct=True,
        # coroutine=aMultiply
    )
    print(calculate.invoke({"a": 2, "b": 3}))
    print(f"名称：{calculate.name}")
    print(f"描述：{calculate.description}")
    print(f"参数：{calculate.args}")

    # print(await calculate.ainvoke({"a": 2, "b": 5}))

# 运行异步主函数
asyncio.run(main())
